import numpy as np  # 用于矩阵操作
import gradio as gr  # 用于创建网页界面
from sklearn.cluster import KMeans  # 用于K-means聚类算法

# ########## 函数：使用K-means聚类提取主色调 ##########
def extract_dominant_colors(image, k=5):
    """
    使用K-means聚类提取图像的主色调
    :param image: 输入图像（RGB格式）
    :param k: 聚类的簇数（默认为5）
    :return: 主色调的RGB值
    """
    # 1. 将图像数据转为一维像素数组
    img = image.copy()
    img = img.reshape((-1, 3))  # 转换为二维矩阵，每行是一个像素点（RGB值）
    img = np.float32(img)  # 转换为float32类型

    # 2. 使用K-means算法进行聚类
    kmeans = KMeans(n_clusters=k, random_state=0, n_init=10)  # 创建K-means模型
    kmeans.fit(img)  # 执行聚类
    centers = kmeans.cluster_centers_  # 获取聚类中心（主色调）
    
    # 3. 获取每个簇的RGB值
    dominant_colors = np.uint8(centers)  # 聚类中心的RGB值
    return dominant_colors

# ########## 函数：展示主色调和十六进制编码 ##########
def display_dominant_colors(image):
    # 1. 读取图像
    original_image = image  # Gradio输入的图像已经是RGB格式

    # 2. 提取主色调
    dominant_colors = extract_dominant_colors(original_image, k=5)

    # 3. 构建一个色板显示主色调
    color_palette = np.zeros((100, 500, 3), dtype=np.uint8)
    step = 500 // 5
    hex_codes = []  # 用于保存十六进制编码
    for i in range(5):
        color_palette[:, i*step:(i+1)*step, :] = dominant_colors[i]
        # 转换RGB到十六进制
        hex_code = '#{:02x}{:02x}{:02x}'.format(dominant_colors[i][0], dominant_colors[i][1], dominant_colors[i][2])
        hex_codes.append(hex_code)

    # 4. 返回色板和十六进制编码
    hex_codes_text = '\n'.join(hex_codes)  # 合并为一个字符串，按行显示
    return color_palette, hex_codes_text

# ########## Gradio界面 ##########
iface = gr.Interface(
    fn=display_dominant_colors,  # 主函数
    inputs=gr.Image(type="numpy"),  # 输入图片（已更新为gr.Image）
    outputs=[gr.Image(label="主色调"), gr.Textbox(label="主色调的十六进制编码")],  # 输出主色调色板和十六进制编码
    live=True,  # 实时更新
    title="图片主色调提取",  # 标题
    description="上传一张图片，提取其主色调（5个）并展示，显示每个主色调的十六进制编码",  # 描述
)

iface.launch()  # 启动Gradio界面
